﻿#pragma once
#include <cstdio>
#include <cstdint>
#include <bits/stl_function.h>
#include <queue>
#include <assert.h>

/// 二叉树节点类
struct BinNodeBase
{
    using Ptr = BinNodeBase*;
    using ConstPtr = const BinNodeBase*;

    Ptr m_parent{ nullptr };
    Ptr m_left{ nullptr };
    Ptr m_right{ nullptr };
};

/// 红黑树节点颜色
enum RBColor
{
    kRBColorRed = false,
    kRBColorBlack = true,
};

/// 判断是否是左孩子
#define IsLeftChild(x) ((x) && ((x) == (x)->m_parent->m_left))
/// 判断节点颜色是否为黑色
#define IsBlack(x) (!(x) || RBColor::kRBColorBlack == (x)->m_color)
/// 判断节点颜色是否为红色
#define IsRed(x) (x && RBColor::kRBColorRed == (x)->m_color)

/// 红黑树节点基类
struct RBNodeBase
{
    using BasePtr = RBNodeBase*;
    using ConstBasePtr = const RBNodeBase*;
    using ColorType = RBColor;

    /// 节点颜色（放在最上面方便debug查看）
    ColorType m_color{ RBColor::kRBColorRed };
    BasePtr m_parent{ nullptr };
    BasePtr m_left{ nullptr };
    BasePtr m_right{ nullptr };
#ifdef DEBUG
    uint32_t m_debug_val{ 0 };
#endif

    static BasePtr
    minimum(BasePtr x)
    {
        while (x && x->m_left != nullptr)
        {
            x = x->m_left;
        }
        return x;
    }

    static BasePtr
    maximum(BasePtr x)
    {
        while (x && x->m_right != nullptr)
        {
            x = x->m_right;
        }
        return x;
    }

    inline static bool
    IsRightChild(BasePtr x)
    {
        return x && x == x->m_parent->m_right;
    }

    /// 叔父
    inline static BasePtr
    uncle(BasePtr x)
    {
        return IsLeftChild(x->m_parent) 
            ? x->m_parent->m_parent->m_right
            : x->m_parent->m_parent->m_left;
    }

    /// 兄弟
    static BasePtr
    sibling(BasePtr x)
    {
        return IsLeftChild(x) 
            ? x->m_parent->m_right
            : x->m_parent->m_left;
    }
};

/// 红黑树节点模板类
template<typename Value>
struct RBNode : public RBNodeBase
{
    using ValueType = Value;

    ValueType m_val;
};

struct RBTreeImpl
{
    RBNodeBase m_header;
    std::size_t m_node_count{ 0 };

    RBTreeImpl()
    {
        m_header.m_color = RBColor::kRBColorRed;
        Reset();
    }

    void Reset()
    {
        m_header.m_parent = nullptr;
        m_header.m_left = &m_header;
        m_header.m_right = &m_header;
        m_node_count = 0u;
    }
};

struct RBTreeIteratorBase
{
};

static RBNodeBase *
RBTreeIncrment(RBNodeBase *x);

static RBNodeBase *
RBTreeDecrement(RBNodeBase *x);

static void
RBTreeInsert(const bool insert_left,
             RBNodeBase *x,
             RBNodeBase *p,
             RBNodeBase &header);

static RBNodeBase*
RBTreeErase(RBNodeBase* const z, 
            RBNodeBase &header);


/// 红黑树迭代器
template<typename Value>
struct RBTreeIterator
{
    using Self = RBTreeIterator<Value>;
    using RBNodeBasePtr = typename RBNodeBase::BasePtr;
    using RBNodePtr = RBNode<Value>*;
    using ConstRBNodeBasePtr = const typename RBNodeBase::BasePtr;
    using ConstRBNodePtr = const RBNode<Value>*;
    using ValueType = Value;
    using RefType = Value&;
    using PointType = Value*;

    RBNodeBasePtr m_node{ nullptr };

    RBTreeIterator()
        : m_node()
    {
    }

    explicit
    RBTreeIterator(RBNodeBasePtr x)
        : m_node(x)
    {
    }

    void
    Increment()
    {
        m_node = RBTreeIncrment(m_node);
    }
    
    void
    Decrement()
    {
        m_node = RBTreeDecrement(m_node);
    }

    RefType
    operator*() const
    {
        return static_cast<RBNodePtr>(m_node)->m_val;
    }

    PointType
    operator->() const
    {
        return &static_cast<RBNodePtr>(m_node)->m_val;
    }

    Self&
    operator++()
    {
        Increment();
        return *this;
    }

    Self
    operator++(int)
    {
        Self tmp = *this;
        Increment();
        return tmp;
    }

    Self&
    operator--()
    {
        Decrement();
        return *this;
    }

    Self
    operator--(int)
    {
        Self tmp = *this;
        Decrement();
        return tmp;
    }

    bool
    operator==(const Self &x) const
    {
        return m_node == x.m_node;
    }

    bool
    operator!=(const Self &x) const
    {
        return m_node != x.m_node;
    }
};

/// 红黑树模板类
template<typename Key, typename Value, 
    typename KeyOfValue = std::_Identity<Value>, 
    typename Compare = std::less<Key>,
    typename Alloc = std::allocator<Value>>
class RBTree
{
public:
    using KeyType = Key;
    using ValueType = Value;
    using AllocType = Alloc;
    using RBNodePtr = RBNode<ValueType>*;
    using ConstRBNodePtr = const RBNodePtr;
    using RBNodeBasePtr = RBNodeBase*;
    using ConstRBNodeBasePtr = const RBNodeBasePtr;
    using Iterator = RBTreeIterator<ValueType>;
    using ConstIterator = const RBTreeIterator<ValueType>;

    using InsertReturnType = std::pair<Iterator, bool>;

    using size_type = std::size_t;

protected:
    RBTreeImpl m_impl;

public:
    //AllocType GetAllocator() const { return AllocType(); }

    RBTree() = default;
    
    ~RBTree() { clear(); }

public:
    Iterator begin() const { return Iterator(m_impl.m_header.m_left); }
    Iterator end() { return Iterator(&m_impl.m_header); }
private:
    RBNodeBasePtr _headerBase() noexcept { return &m_impl.m_header; }
    RBNodeBasePtr _rootBase() noexcept { return m_impl.m_header.m_parent; }
    RBNodeBasePtr _leftMostBase() noexcept { return m_impl.m_header.m_left; }
    RBNodeBasePtr _rightMostBase() noexcept { return m_impl.m_header.m_right; }
    RBNodeBasePtr _beginBase() noexcept { return m_impl.m_header.m_left; }
    RBNodeBasePtr _endBase() noexcept { return &m_impl.m_header; }
    
    RBNodePtr _root() noexcept { return static_cast<RBNodePtr>(m_impl.m_header.m_parent); }
    
    //RBNodeBase& _header() noexcept { return m_impl.m_header; }
    //ConstRBNodePtr _root() const noexcept { return static_cast<ConstRBNodePtr>(m_impl.m_header.m_parent); }
    //RBNodeBasePtr _endBase() noexcept { return static_cast<RBNodeBasePtr>(&m_impl.m_header); }
    //ConstRBNodePtr _left(ConstRBNodePtr x) const noexcept { return static_cast<ConstRBNodePtr>(x->m_left); }
    //ConstRBNodePtr _right(ConstRBNodePtr x) const noexcept { return static_cast<ConstRBNodePtr>(x->m_right); }
private:
    static RBNodePtr
    _left(RBNodePtr x) noexcept
    { return static_cast<RBNodePtr>(x->m_left); }
    
    static RBNodePtr
    _right(RBNodePtr x) noexcept
    { return static_cast<RBNodePtr>(x->m_right); }
public:
    size_type
    size() const
    {
        return m_impl.m_node_count;
    }

    void clear()
    {
        _erase(_root());
        m_impl.Reset();
    }

    InsertReturnType
    insertUnique(const ValueType &v)
    {
        using RetType = InsertReturnType;
        std::pair<RBNodeBasePtr, RBNodeBasePtr> res = getInsertUniquePos(v);
        if (res.second)
        {
            return RetType(_insert(res.first, res.second, v),
                    true);
        }
        return RetType(Iterator(res.first), false);
    }

    InsertReturnType
    insertEqual(const ValueType &v)
    {
        using RetType = InsertReturnType;
        return RetType(Iterator(nullptr), false);
    }

    ConstIterator
    find(const KeyType &k)
    {
        Iterator i = lowerBound(_root(), _headerBase(), k);
        return (i == end() || k < *i) ? end() : i;
    }

    std::pair<ConstIterator, ConstIterator>
    equalRange(const ValueType &k)
    {// stl_tree.h : 2024
        using RetType = std::pair<ConstIterator, ConstIterator>;
        RBNodeBasePtr x = _rootBase();
        RBNodeBasePtr y = _endBase();
        while (x != nullptr)
        {
            if (static_cast<RBNodePtr>(x)->m_val < k)
            {// x < k
                x = x->m_right;
            }
            else if (k < static_cast<RBNodePtr>(x)->m_val)
            {// k < x
                y = x;
                x = x->m_left;
            }
            else
            {// x == k
                // TODO
                //return RetType(lowerBound(x, y, k), upperBound(x, y, k));
                return RetType(ConstIterator(x), ConstIterator(y));
            }
        }
        return RetType(ConstIterator(y), ConstIterator(y));
    }

    Iterator
    lowerBound(const KeyType &k) const
    {
        return Iterator();
    }

    size_type
    erase(const KeyType &k)
    {//stl_tree.h : 2534
        std::pair<ConstIterator, ConstIterator> p = equalRange(k);
        if (p.first == end())
        {
            return 0;
        }
        const size_type old_size = size();
        {
            RBNodeBasePtr x = RBTreeErase(p.first.m_node, m_impl.m_header);
            putNode(static_cast<RBNodePtr>(x));
        }
        m_impl.m_node_count--;
        return old_size - size();
    }
protected:
    RBNodePtr
    getNode()
    {
        // TODO
        return new RBNode<ValueType>();
    }

    void
    putNode(RBNodePtr x)
    {
        delete x;
    }
private:
    Iterator
    lowerBound(RBNodePtr x, RBNodeBasePtr p, const ValueType &k) const;

    Iterator
    upperBound(RBNodePtr x, RBNodeBasePtr y, const ValueType &k)
    {
        while (x != nullptr)
        {
            if (k < x->m_val)
            {// x > k
                y = x;
                x = _left(x);
            }
            else
            {// x <= k
                x = _right(x);
            }
        }
        return Iterator(y);
    }

    std::pair<RBNodeBasePtr, RBNodeBasePtr>
    getInsertUniquePos(const ValueType &v)
    {
        using RetType = std::pair<RBNodeBasePtr, RBNodeBasePtr>;
        RBNodeBasePtr y = _headerBase();
        RBNodeBasePtr x = _rootBase();
        bool insert_left = true;
        while (x != nullptr)
        {
            y = x;
            insert_left = v < static_cast<RBNodePtr>(x)->m_val;
            x = insert_left ? x->m_left : x->m_right;
        }
        Iterator i = Iterator(y);
        if (insert_left)
        {
            if (i == begin())
            {
                return RetType(x, y);
            }
            else
            {// 向上回到比它小的第一个祖先
                --i;
            }
        }
        // 再次比较大小
        // 由于处理过向左插入的情况，此时i一定小于v
        // 否则出现重复元素，插入失败
        if (static_cast<RBNodePtr>(i.m_node)->m_val < v)
        {
            return RetType(x, y);
        }
        else
        {
            return RetType(i.m_node, nullptr);
        }
    }

    RBNodePtr createNode(const ValueType &v)
    {
        //RBNodePtr x = getNode();
        RBNodePtr x = new RBNode<ValueType>();
        // TODO
        //x->m_val = new ValueType(v);
        x->m_val = v;
        return x;
    }

    void destoryNode(RBNodePtr x)
    {
        //delete x->m_val;
        putNode(x);
    }
public:
    
    bool remove(const ValueType &x);

    /// 插入元素
    Iterator _insert(RBNodeBasePtr x, RBNodeBasePtr p, const ValueType &v)
    {
        bool insert_left = (p == _endBase()
                            || v < static_cast<RBNodePtr>(p)->m_val
                            || nullptr != x);
        RBNodePtr z = createNode(v);
#ifdef DEBUG
        z->m_debug_val = v;
#endif
        RBTreeInsert(insert_left, z, p, m_impl.m_header);
        ++m_impl.m_node_count;
        return Iterator(z);
    }

    /// 删除子树无调整
    void _erase(RBNodePtr x)
    {
        RBNodePtr y = nullptr;
        while (x != nullptr)
        {
            _erase(_right(x));
            y = _left(x);
            destoryNode(x);
            x = y;
        }
    }

    void _debug()
    {
        int32_t cnt = 0;
        int32_t nxt_cnt = 1;
        std::queue<RBNodeBasePtr> que;
        que.emplace(_rootBase());
        while (!que.empty())
        {
            RBNodeBasePtr u = que.front();
            que.pop();
            if (++cnt == nxt_cnt
                || u == _rootBase()
                || (u && RBNodeBase::IsRightChild(u->m_parent))
                || (u && (RBNodeBase::IsRightChild(u) && u->m_parent == _rootBase())))
            {
                cnt = 0;
                nxt_cnt <<= 1;
                printf("\n");
            }
            if (u)
            {
                que.emplace(u->m_left);
                que.emplace(u->m_right);
            }
        }
        printf("\n");
    }
};

template<typename Key, typename Value, 
    typename KeyOfValue, typename Compare, typename Alloc>
typename RBTree<Key, Value, KeyOfValue, Compare, Alloc>::Iterator
RBTree<Key, Value, KeyOfValue, Compare, Alloc>::
lowerBound(RBNodePtr x, RBNodeBasePtr p, const ValueType &k) const
{
    while (x != nullptr)
    {
        if (!(x->m_val < k))
        {// x >= k
            p = x;
            x = _left(x);
        }
        else
        {// x < k
            x = _right(x);
        }
    }
    return Iterator(p);
}

#include "rb_tree_inl.h"